{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import os\n",
    "import functools\n",
    "\n",
    "import cV3\n",
    "import numpy as np\n",
    "\n",
    "from tests.metaworld.envs.mujoco.sawyer_xyz.test_scripted_policies import ALL_ENVS, test_cases_latest_nonoise\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "def trajectory_generator(env, policy, act_noise_pct, res=(640, 480), camera='corner'):\n",
    "    action_space_ptp = env.action_space.high - env.action_space.low\n",
    "\n",
    "    env.reset()\n",
    "    env.reset_model()\n",
    "    o = env.reset()\n",
    "\n",
    "    for _ in range(env.max_path_length):\n",
    "        a = policy.get_action(o)\n",
    "        a = np.random.normal(a, act_noise_pct * action_space_ptp)\n",
    "\n",
    "        o, r, done, info = env.step(a)\n",
    "        # Camera is one of ['corner', 'topview', 'behindGripper', 'gripperPOV']\n",
    "        yield r, done, info, env.sim.render(*res, mode='offscreen', camera_name=camera)[:,:,::-1]\n",
    "\n",
    "def writer_for(tag, fps, res):\n",
    "    if not os.path.exists('movies'):\n",
    "        os.mkdir('movies')\n",
    "    return cV3.VideoWriter(\n",
    "        f'movies/{tag}.avi',\n",
    "        cV3.VideoWriter_fourcc('M','J','P','G'),\n",
    "        fps,\n",
    "        res\n",
    "    )"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "resolution = (1920, 1080)\n",
    "camera = 'behindGripper' # one of ['corner', 'topview', 'behindGripper', 'gripperPOV']\n",
    "flip=True # if True, flips output image 180 degrees\n",
    "\n",
    "config = [\n",
    "    # env, action noise pct, cycles, quit on success\n",
    "    ('assembly-V3', np.zeros(4), 3, True),\n",
    "    ('basketball-V3', np.zeros(4), 3, True),\n",
    "    ('bin-picking-V3', np.zeros(4), 3, True),\n",
    "    ('box-close-V3', np.zeros(4), 3, True),\n",
    "    ('button-press-topdown-V3', np.zeros(4), 3, True),\n",
    "    ('button-press-topdown-wall-V3', np.zeros(4), 3, True),\n",
    "    ('button-press-V3', np.zeros(4), 3, True),\n",
    "    ('button-press-wall-V3', np.zeros(4), 3, True),\n",
    "    ('coffee-button-V3', np.zeros(4), 3, True),\n",
    "    ('coffee-pull-V3', np.zeros(4), 3, True),\n",
    "    ('coffee-push-V3', np.zeros(4), 3, True),\n",
    "    ('dial-turn-V3', np.zeros(4), 3, True),\n",
    "    ('disassemble-V3', np.zeros(4), 3, True),\n",
    "    ('door-close-V3', np.zeros(4), 3, True),\n",
    "    ('door-lock-V3', np.zeros(4), 3, True),\n",
    "    ('door-open-V3', np.zeros(4), 3, True),\n",
    "    ('door-unlock-V3', np.zeros(4), 3, True),\n",
    "    ('hand-insert-V3', np.zeros(4), 3, True),\n",
    "    ('drawer-close-V3', np.zeros(4), 3, True),\n",
    "    ('drawer-open-V3', np.zeros(4), 3, True),\n",
    "    ('faucet-open-V3', np.zeros(4), 3, True),\n",
    "    ('faucet-close-V3', np.zeros(4), 3, True),\n",
    "    ('hammer-V3', np.zeros(4), 3, True),\n",
    "    ('handle-press-side-V3', np.zeros(4), 3, True),\n",
    "    ('handle-press-V3', np.zeros(4), 3, True),\n",
    "    ('handle-pull-side-V3', np.zeros(4), 3, True),\n",
    "    ('handle-pull-V3', np.zeros(4), 3, True),\n",
    "    ('lever-pull-V3', np.zeros(4), 3, True),\n",
    "    ('peg-insert-side-V3', np.zeros(4), 3, True),\n",
    "    ('pick-place-wall-V3', np.zeros(4), 3, True),\n",
    "    ('pick-out-of-hole-V3', np.zeros(4), 3, True),\n",
    "    ('reach-V3', np.zeros(4), 3, True),\n",
    "    ('push-back-V3', np.zeros(4), 3, True),\n",
    "    ('push-V3', np.zeros(4), 3, True),\n",
    "    ('pick-place-V3', np.zeros(4), 3, True),\n",
    "    ('plate-slide-V3', np.zeros(4), 3, True),\n",
    "    ('plate-slide-side-V3', np.zeros(4), 3, True),\n",
    "    ('plate-slide-back-V3', np.zeros(4), 3, True),\n",
    "    ('plate-slide-back-side-V3', np.zeros(4), 3, True),\n",
    "    ('peg-insert-side-V3', np.zeros(4), 3, True),\n",
    "    ('peg-unplug-side-V3', np.zeros(4), 3, True),\n",
    "    ('soccer-V3', np.zeros(4), 3, True),\n",
    "    ('stick-push-V3', np.zeros(4), 3, True),\n",
    "    ('stick-pull-V3', np.zeros(4), 3, True),\n",
    "    ('push-wall-V3', np.zeros(4), 3, True),\n",
    "    ('push-V3', np.zeros(4), 3, True),\n",
    "    ('reach-wall-V3', np.zeros(4), 3, True),\n",
    "    ('reach-V3', np.zeros(4), 3, True),\n",
    "    ('shelf-place-V3', np.zeros(4), 3, True),\n",
    "    ('sweep-into-V3', np.zeros(4), 3, True),\n",
    "    ('sweep-V3', np.zeros(4), 3, True),\n",
    "    ('window-open-V3', np.zeros(4), 3, True),\n",
    "    ('window-close-V3', np.zeros(4), 3, True),\n",
    "]\n",
    "\n",
    "for env, noise, cycles, quit_on_success in config:\n",
    "    tag = env + '-noise-' + np.array2string(noise, precision=2, separator=',', suppress_small=True)\n",
    "\n",
    "    policy = functools.reduce(lambda a,b : a if a[0] == env else b, test_cases_latest_nonoise)[1]\n",
    "    env = ALL_ENVS[env]()\n",
    "    env._partially_observable = False\n",
    "    env._freeze_rand_vec = False\n",
    "    env._set_task_called = True\n",
    "\n",
    "    writer = writer_for(tag, env.metadata['video.frames_per_second'], resolution)\n",
    "    for _ in range(cycles):\n",
    "        for r, done, info, img in trajectory_generator(env, policy, noise, resolution, camera):\n",
    "            if flip: img = cV3.rotate(img, cV3.ROTATE_180)\n",
    "            writer.write(img)\n",
    "            if quit_on_success and info['success']:\n",
    "                break\n",
    "\n",
    "    writer.release()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
